import sys
from pathlib import Path

file = Path(__file__).resolve()
parent, top = file.parent, file.parents[2]

sys.path.append(str(top))
try:
    sys.path.remove(str(parent))
except ValueError: # Already removed
    pass

__package__ = 'src.job'


from jd_package_oqc.jdoqc_interface import test_full_image_search_oqc_folder_panels
from ..config import material_dir, lot_dir, thumb_dir, lot_upload_dir
from ..utils.log import logger
from .model import det_model_1, det_model_2, det_model_3, cls_model_forWhite, cls_model_forRules
import os
from ..utils.database import WebSessionLocal, InferSessionLocal
from ..model.infer import InferPanel, InferPanelImage, Failure, Material, Lot
from ..model.web import Panel, PanelImage, Failure, Workstation, PanelWorkstation
import json
from ..utils.file import image_thumb
from datetime import datetime, timedelta
from .. import async_msg


class PanelJob(object):
    """
    Base class
    """
    def __init__(self):
        self.det_models_dict = {"det_model_1": det_model_1, "det_model_2": det_model_2, "det_model_3": det_model_3}
        self.cls_models_dict = {"cls_model_forWhite": cls_model_forWhite, "cls_model_forRules": cls_model_forRules}


    def detect_lot(self, lot_path, material_code, lot_code=None, aoi_code=None, is_last=None):
        # 写入 infer 数据库, 写入 web 数据库
        with InferSessionLocal() as infer_session, WebSessionLocal() as web_session:
            try:
                material_obj = infer_session.query(Material).filter_by(material_code=material_code).first()
                material_name = material_obj.material_name
                ini_path = f"{material_dir}/{material_name}/parameter.ini"

                if lot_code:
                    logger.info("开始执行自动推理任务")
                    # 自动拷贝
                    lot = infer_session.query(Lot).filter_by(lot_code=lot_code).first()
                    # 自动生成的 lot 名
                    lotname = os.path.basename(lot.dir_path)
                    # 人工输入的 lot 名
                    lot_name = lot.lot_name
                    lot_folder = lot_dir

                else: 
                    msg = "开始执行手动推理任务"
                    logger.info(msg)

                    # 手动上传
                    lotname = os.path.basename(lot_path)

                    lot_name = lotname

                    # 数据库中存储的是 hdd 上的 lot 路径
                    lot = Lot(lot_name=lotname, dir_path=lot_path, material_code=material_code)
                    infer_session.add(lot)
                    infer_session.commit()
                    lot_code = lot.lot_code
                    lot_folder = lot_upload_dir

                # lot_folder: 存放 lot 的目录
                thumb_folder = lot_path.replace(lot_folder, thumb_dir)

                panels_result_oqc_list, panels_aoi_result_oqc_list = test_full_image_search_oqc_folder_panels(
                    lot_path,
                    self.det_models_dict,   
                    self.cls_models_dict,   
                    ini_path
                )

                for i, single_panel_ai_list in enumerate(panels_result_oqc_list):
                    panel_id = single_panel_ai_list["panel_id"]
                    ai_results = single_panel_ai_list["result"]
                    img_roi = single_panel_ai_list["img_roi"]
                    ogn = single_panel_ai_list["OGN"]
                    img_status = single_panel_ai_list["img_status"]                        

                    single_panel_aoi_list = panels_aoi_result_oqc_list[i]
                    aoi_results = single_panel_aoi_list["result"]

                    thumb_file_folder = os.path.join(thumb_folder, panel_id)
                    os.makedirs(thumb_file_folder, exist_ok=True)

                    panel_id_info = panel_id.split("_")

                    date_str = panel_id_info[2]
                    current_year = datetime.now().year
                    datetime_str = f"{current_year}{date_str[:2]}{date_str[2:4]}{date_str[4:6]}{date_str[6:8]}{date_str[8:10]}"

                    panel_time= datetime.strptime(datetime_str, "%Y%m%d%H%M%S")

                    if panel_id_info[0]=="NG":
                        aoi_result = 1 
                    elif panel_id_info[0]=="OK":
                        aoi_result = 0                   

                    # 先写入 infer 表获取 panel code
                    infer_panel = InferPanel(
                        lot_code=lot_code,
                        material_code=material_code,
                        aoi_code=aoi_code,
                        panel_id=panel_id,
                        aoi_result=aoi_result,
                        ai_result=ogn,
                        start_time=panel_time,
                        img_roi=json.dumps(img_roi)
                    )

                    infer_session.add(infer_panel)
                    infer_session.commit()

                    panel_code = infer_panel.panel_code

                    # 写入 web 表
                    panel = Panel(panel_code=panel_code, lot_name=lot_name, bar_code=panel_id, ai=ogn, start_time=panel_time, aoi=aoi_result, material_name=material_name, img_roi=json.dumps(img_roi))

                    web_session.add(panel)
                    web_session.commit()

                    workstations = web_session.query(Workstation).filter_by(del_flag=0).all()

                    # 写入 web 表
                    for workstation in workstations:
                        panel_workstation = PanelWorkstation(panel_code=panel.panel_code, workstation_code=workstation.id)
                        web_session.add(panel_workstation)
                        web_session.commit()

                    for image_name, ai_failures in ai_results.items():
                        panel_image_path = os.path.join(lot_path, panel_id, "photo", image_name)
                        thumb_file_path = os.path.join(thumb_file_folder, image_name)

                        try:
                            image_thumb(panel_image_path, thumb_file_path)
                        except Exception as e:
                            logger.error(f"缩略图 {panel_id} 生成失败: {e}")
                            thumb_file_path = ""

                        if img_status == 'bad':
                            ai_detection_info = json.dumps({"ZA11": [{"box": [0, 0, 100, 100], "score": 1}]})
                        else:
                            ai_detection_info = json.dumps(ai_failures) if ai_failures else ""

                        aoi_detection_info = json.dumps(aoi_results[image_name]) if aoi_results[image_name] else ""

                        pattern_failure = image_name.split(".")[0]

                        # 写入 infer 表
                        infer_panel_image = InferPanelImage(panel_code=panel_code, lot_code=lot_code, image_path=panel_image_path, detection_info=ai_detection_info, image_name=image_name, thumb_file_path=thumb_file_path, aoi_failures=aoi_detection_info, pattern_failure=pattern_failure)
                        infer_session.add(infer_panel_image)
                        infer_session.commit()

                        # 写入 web 表
                        # image_path 使用相对路径
                        web_panel_image_path = os.path.join(panel_id, "photo", image_name)
                        panel_image = PanelImage(panel_code=panel_code, detection_info=ai_detection_info, image_name=image_name, image_path=web_panel_image_path, aoi_failures=aoi_detection_info, pattern_failure=pattern_failure)
                        web_session.add(panel_image)
                        web_session.commit()

                    web_session.commit()

                # 如果是手动上传
                if not aoi_code:
                    lot.status = "1"
                lot.end_time = datetime.utcnow() + timedelta(hours=8)
                infer_session.commit()
            except Exception as e:
                logger.error(f"推理出错: {e}")

                if lot:
                    lot.status = '2'
                    lot.end_time = datetime.utcnow() + timedelta(hours=8)
                    infer_session.commit()


panel_job = PanelJob()

if __name__ == "__main__":
    lot_folder = "/nfs/home/xh/OQC/tiny_Data"

    panel_job.detect_folder(lot_folder)


